package com.jd.genie.agent.tool.common;

import com.alibaba.fastjson.JSON;
import com.jd.genie.agent.agent.AgentContext;
import com.jd.genie.agent.dto.Plan;
import com.jd.genie.agent.tool.BaseTool;
import lombok.Data;

import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.function.Function;


/**
 * 计划工具类
 */
@Data
public class PlanningTool implements BaseTool {

    private AgentContext agentContext;
    private final Map<String, Function<Map<String, Object>, String>> commandHandlers = new HashMap<>();
    private Plan plan;

    public PlanningTool()
    {
        commandHandlers.put("create", this::createPlan);
        commandHandlers.put("update", this::updatePlan);
        commandHandlers.put("mark_step", this::markStep);
        commandHandlers.put("finish", this::finishPlan);
    }

    @Override
    public String getName()
    {
        return "planning";
    }

    @Override
    public String getDescription()
    {
        String desc = """
                这是一个计划工具，可让代理创建和管理用于解决复杂任务的计划。
                该工具提供创建计划、更新计划步骤和跟踪进度的功能。
                使用中文回答""";
        return """
                这是一个计划工具，可让代理创建和管理用于解决复杂任务的计划。
                该工具提供创建计划、更新计划步骤和跟踪进度的功能。
                
                创建计划时，需要创建出有依赖关系的计划，计划列表格式如下：
                [
                    执行顺序+编号、任务短标题：任务的细节描述
                ]，
                
                样式示例如下：
                
                [
                    "执行顺序1. 任务短标题: 任务描述xxx ...",
                    "执行顺序1. 任务短标题: 任务描述xxx ...",
                    "执行顺序2. 任务短标题：任务描述xxx ...",
                    "执行顺序3. 任务短标题：任务描述xxx ... "
                ]
                """;
    }

    @Override
    public Map<String, Object> toParams()
    {
        String jsonStr = """
                {
                  "type" : "object",
                  "properties" : {
                    "step_status" : {
                      "description" : "每一个子任务的状态. 当command是 mark_step 时使用.",
                      "type" : "string",
                      "enum" : [ "not_started", "in_progress", "completed", "blocked" ]
                    },
                    "step_notes" : {
                      "description" : "每一个子任务的的备注，当command 是 mark_step 时，是备选参数。",
                      "type" : "string"
                    },
                    "step_index" : {
                      "description" : "当command 是 mark_step 时，是必填参数.",
                      "type" : "integer"
                    },
                    "title" : {
                      "description" : "任务的标题，当command是create时，是必填参数，如果是update 则是选填参数。",
                      "type" : "string"
                    },
                    "steps" : {
                      "description" : "入参是任务列表. 当创建任务时，command是create，此时这个参数是必填参数。任务列表的的格式如下：[\\"执行顺序 + 编号、执行任务简称：执行任务的细节描述\\"]。不同的子任务之间不能重复、也不能交叠，可以收集多个方面的信息，收集信息、查询数据等此类多次工具调用，是可以并行的任务。具体的格式示例如下：- 任务列表示例1: [\\"执行顺序1. 执行任务简称（不超过6个字）：执行任务的细节描述（不超过50个字）\\", \\"执行顺序2. xxx（不超过6个字）：xxx（不超过50个字）, ...\\"]；",
                      "type" : "array",
                      "items" : {
                        "type" : "string"
                      }
                    },
                    "command" : {
                      "description" : "需要执行的命令，取值范围是: create, update, mark_step",
                      "type" : "string",
                      "enum" : [ "create", "update", "mark_step" ]
                    }
                  },
                  "required" : [ "command" ]
                }
                """;
        return JSON.parseObject(jsonStr, Map.class);
    }

    @Override
    public Object execute(Object input)
    {
        if (!(input instanceof Map)) {
            throw new IllegalArgumentException("Input must be a Map");
        }

        Map<String, Object> params = (Map<String, Object>) input;
        String command = (String) params.get("command");
        if (command == null || command.isEmpty()) {
            throw new IllegalArgumentException("Command is required");
        }

        Function<Map<String, Object>, String> handler = commandHandlers.get(command);
        if (handler != null) {
            return handler.apply(params);
        } else {
            throw new IllegalArgumentException("Unknown command: " + command);
        }
    }

    private String createPlan(Map<String, Object> params)
    {
        String title = (String) params.get("title");
        List<String> steps = (List<String>) params.get("steps");

        if (title == null || steps == null) {
            throw new IllegalArgumentException("title, and steps are required for create command");
        }

        if (plan != null) {
            throw new IllegalStateException("A plan already exists. Delete the current plan first.");
        }

        plan = Plan.create(title, steps);
        return "我已创建plan";
    }

    private String updatePlan(Map<String, Object> params)
    {
        String title = (String) params.get("title");
        List<String> steps = (List<String>) params.get("steps");

        if (plan == null) {
            throw new IllegalStateException("No plan exists. Create a plan first.");
        }

        plan.update(title, steps);
        return "我已更新plan";
    }

    private String markStep(Map<String, Object> params)
    {
        Integer stepIndex = (Integer) params.get("step_index");
        String stepStatus = (String) params.get("step_status");
        String stepNotes = (String) params.get("step_notes");

        if (plan == null) {
            throw new IllegalStateException("No plan exists. Create a plan first.");
        }

        if (stepIndex == null) {
            throw new IllegalArgumentException("step_index is required for mark_step command");
        }

        plan.updateStepStatus(stepIndex, stepStatus, stepNotes);

        return String.format("我已标记plan %d 为 %s", stepIndex, stepStatus);
    }

    private String finishPlan(Map<String, Object> params)
    {
        if (Objects.isNull(plan)) {
            plan = new Plan();
        } else {
            for (int stepIndex = 0; stepIndex < plan.getSteps().size(); stepIndex++) {
                plan.updateStepStatus(stepIndex, "completed", "");
            }
        }
        return "我已更新plan为完成状态";
    }

    public void stepPlan()
    {
        plan.stepPlan();
    }

}


